import os

from langchain import hub
from langchain.agents import AgentExecutor
from langchain.agents.format_scratchpad import format_to_openai_function_messages
from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from langchain_core.utils.function_calling import convert_to_openai_function
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.tools import DuckDuckGoSearchResults
from langchain.agents import tool

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "playground"
os.environ["LANGCHAIN_API_KEY"] = "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"


prompt = hub.pull("whatdoing/you_are_doing")
llm = ChatGoogleGenerativeAI(model="models/gemini-1.5-pro-latest", temperature=0.7)
def test1():
    # tools = [get_word_length]
    tools = [
        DuckDuckGoSearchResults(
            name="duck_duck_go"
        ),  # 使用 DuckDuckGo 进行通用互联网搜索
    ]
    llm_with_tools = llm.bind(functions=[convert_to_openai_function(t) for t in tools])
    runnable_agent = (
            {
                "input": lambda x: x["input"],
                "name": lambda x: x["name"],
            }
            | prompt
            | llm_with_tools
            | OpenAIFunctionsAgentOutputParser()
    )

    agent_executor = AgentExecutor(
        agent=runnable_agent, tools=tools, handle_parsing_errors=True
    )

    result = agent_executor.invoke({"input": "小时候", "name":"小明"})
    print(result["output"])

    # results = agent_executor.batch([{"input": x} for x in inputs], return_exceptions=True)
    # print(results)


def get_word_length(word: str) -> int:
    """Returns the length of a word."""
    return len(word)

def test2():
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "You are very powerful assistant, but bad at calculating lengths of words.",
            ),
            ("user", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
    )
    tools = [get_word_length]
    llm_with_tools = llm.bind(functions=[convert_to_openai_function(t) for t in tools])
    agent = (
            {
                "input": lambda x: x["input"],
                "agent_scratchpad": lambda x: format_to_openai_function_messages(
                    x["intermediate_steps"]
                ),
            }
            | prompt
            | llm_with_tools
            | OpenAIFunctionsAgentOutputParser()
    )

    result = agent.invoke({"input": "how many letters in the word educa?", "intermediate_steps": []})
    print(result)
if __name__ == "__main__":
    test1()